def assemble(assembly_code):
    register_map = {
        'ACC': 0, 'BR': 1, 'MR': 2, 'AX': 3,
        'BX': 4, 'CX': 5, 'DX': 6, 'SP': 7
    }

    opcodes = {
        'LOAD': [{'type': 'direct_address', 'opcode': 1}],
        'STORE': [{'type': 'direct_address', 'opcode': 2}],
        'MOV': [{'type': 'reg_reg', 'opcode': 3}, {'type': 'reg_imm', 'opcode': 4}],
        'ADD': [{'type': 'reg', 'opcode': 5}, {'type': 'immediate', 'opcode': 6}],
        'SUB': [{'type': 'reg', 'opcode': 7}, {'type': 'immediate', 'opcode': 8}],
        'MUL': [{'type': 'reg', 'opcode': 9}, {'type': 'immediate', 'opcode': 10}],
        'AND': [{'type': 'reg', 'opcode': 11}, {'type': 'immediate', 'opcode': 12}],
        'OR': [{'type': 'reg', 'opcode': 13}, {'type': 'immediate', 'opcode': 14}],
        'NOT': [{'type': 'none', 'opcode': 15}],
        'SHL': [{'type': 'reg', 'opcode': 16}, {'type': 'immediate', 'opcode': 17}],
        'SHR': [{'type': 'reg', 'opcode': 18}, {'type': 'immediate', 'opcode': 19}],
        'CMP': [{'type': 'reg', 'opcode': 20}, {'type': 'immediate', 'opcode': 21}],
        'JMP': [{'type': 'direct_address', 'opcode': 22}],
        'JZ': [{'type': 'direct_address', 'opcode': 23}],
        'JNZ': [{'type': 'direct_address', 'opcode': 24}],
        'JG': [{'type': 'direct_address', 'opcode': 25}],
        'JNG': [{'type': 'direct_address', 'opcode': 26}],
        'CALL': [{'type': 'direct_address', 'opcode': 27}],
        'RET': [{'type': 'none', 'opcode': 28}],
        'HALT': [{'type': 'none', 'opcode': 29}],
        'PUSH': [{'type': 'reg', 'opcode': 30}],
        'POP': [{'type': 'reg', 'opcode': 31}]
    }

    machine_code = []

    for line in assembly_code.split('\n'):
        line = line.strip()
        if not line or line.startswith(';'):
            continue
        index = line.find(";")
        line = line[:index] if index != -1 else line
        parts = line.split()
        instruction = parts[0].upper()
        operands = ' '.join(parts[1:]) if len(parts) > 1 else ''

        determined_type = None
        if instruction == 'MOV' and ',' in operands:
            split_ops = [op.strip().upper() for op in operands.split(',', 1)]
            if len(split_ops) == 2:
                if split_ops[1].startswith('#'):
                    determined_type = 'reg_imm'
                elif all(op in register_map for op in split_ops):
                    determined_type = 'reg_reg'
                else:
                    raise ValueError(f"Invalid MOV operands: {operands}")
        if not determined_type:
            if operands.startswith('#'):
                determined_type = 'immediate'
            elif ',' in operands:
                split_ops = [op.strip().upper() for op in operands.split(',')]
                if len(split_ops) == 2 and all(op in register_map for op in split_ops):
                    determined_type = 'reg_reg'
                else:
                    raise ValueError(f"Invalid operands for {instruction}: {operands}")
            else:
                if operands.strip().upper() in register_map:
                    determined_type = 'reg'
                else:
                    determined_type = 'direct_address' if operands else 'none'

        opcode_info = next((entry for entry in opcodes.get(instruction, []) if entry['type'] == determined_type), None)
        if not opcode_info:
            raise ValueError(f"Instruction {instruction} does not support operand type {determined_type}")

        opcode = opcode_info['opcode']
        instr_type = opcode_info['type']
        opcode_bin = format(opcode, '06b')
        operand_bin = ''

        if instr_type == 'direct_address':
            addr = int(operands)
            if not 0 <= addr <= 255:
                raise ValueError(f"Direct address {addr} out of range (0-255)")
            operand_bin = format(addr, '08b').zfill(10)
        elif instr_type == 'reg_reg':
            rd, rs = [op.strip().upper() for op in operands.split(',')]
            rd_num = register_map[rd]
            rs_num = register_map[rs]
            operand_bin = f"{rd_num:03b}{rs_num:03b}0000"
        elif instr_type == 'reg':
            rd = operands.strip().upper()
            rd_num = register_map[rd]
            if instruction == 'POP' :
                operand_bin = f"{rd_num:03b}0000000"#单寄存器只有POP是输入，其他都是输出[9:7]输入 [6:4]输出
            else:
                operand_bin = f"000{rd_num:03b}0000"
        elif instr_type == 'immediate':
            imm = int(operands[1:]) if operands[1] != '-' else int(operands[1:])
            if not -512 <= imm <= 511:
                raise ValueError(f"Immediate value {imm} out of range (-512 to 511)")
            operand_bin = format(imm & 0x3FF, '010b')  # 10-bit 2’s complement
        elif instr_type == 'reg_imm':  
            rd_str, imm_str = [op.strip().upper() for op in operands.split(',', 1)]
            rd_num = register_map.get(rd_str, -1)
            imm = int(imm_str[1:]) if imm_str[1] != '-' else int(imm_str[1:])
            if rd_num == -1:
                raise ValueError(f"Invalid register: {rd_str}")
            if not -64 <= imm <= 63:
                raise ValueError(f"Immediate value {imm} out of range (-64 to 63)")
            operand_bin = f"{rd_num:03b}{format(imm & 0x7F, '07b')}"  # 7-bit 2’s complement
        elif instr_type == 'none':
            operand_bin = '0' * 10

        full_instruction = opcode_bin + operand_bin
        machine_code.append(full_instruction)

    return machine_code

# 示例
assembly = """
MOV SP,#20
MOV BX,#0
MOV AX,#1
MOV ACC,BX
ADD AX
MOV BX,ACC
MOV ACC,AX
ADD #1
MOV AX,ACC
CMP #101
JNG 3
MOV ACC,BX
HALT        ; 直接寻址
"""

machine_code = assemble(assembly)
print(machine_code)

with open('instruction.coe', 'w', newline='', encoding='utf-8') as outfile_coe:
    outfile_coe.write("MEMORY_INITIALIZATION_RADIX=2;\n")
    outfile_coe.write("MEMORY_INITIALIZATION_VECTOR=\n")
    for i, row in enumerate(machine_code):
        outfile_coe.write(row)
        if i < len(machine_code) - 1:
            outfile_coe.write(",\n")
        else:
            outfile_coe.write(";\n")


# MOV ACC,#-2
# MUL #18
# MOV AX, #2
# MOV CX, #-9
# MOV ACC, AX
# MUL #3
# MUL CX
# HALT